#导入必要的库
import cv2
import numpy as np
import os
import time

##加载训练数据
def load_data():
    train_data = []
    train_labels = []
    for fi in range(1, 37):
        img=cv2.imread('x\\'+str(fi)+'.png')
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        train_data.append(img)
        train_labels.append(int('0000000001221000000001221000000001221'[fi]))
    train_data = np.array(train_data)
    train_labels = np.array(train_labels)
    train_data = train_data.astype(np.float32)
    return train_data, train_labels

def train():
    ##训练模型
    train_data, train_labels = load_data()
    # 将数据转换为一维数组
    train_data = train_data.reshape((train_data.shape[0], -1))
    # 训练模型
    model = cv2.ml.SVM_create()
    model.setType(cv2.ml.SVM_C_SVC)
    model.setKernel(cv2.ml.SVM_LINEAR)
    model.setTermCriteria((cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-6))
    result = model.train(train_data, cv2.ml.ROW_SAMPLE, train_labels)#cv2.ml.ROW_SAMPLE, 
    model.save('go.xml')
    return model

def test(model):
    #for fi in range(1, 8):
    #f='k\\'+str(fi)+'.png'
    f='k\\1.png'
    _img=cv2.imread(f)
    _img = cv2.cvtColor(_img, cv2.COLOR_BGR2GRAY)
    test_data=[]
    for py in range(19):
        y=py*32+1
        for px in range(19):
            x=px*32+1
            cimg=_img[y:y+32, x:x+32]
            test_data.append(cimg)
    _test(model, test_data)

def _test(model, test_data):
    test_data = np.array(test_data)
    test_data = test_data.astype(np.float32)
    test_data = test_data.reshape((test_data.shape[0], -1))
    # 进行预测
    _, test_labels = model.predict(test_data)
    res=test_labels.reshape((19, 19))
    s=res[3][3]
    print(s)
    if(s==2):
        print('ok')
    print(res)

if __name__ == '__main__':
    train()
    model = cv2.ml.SVM_load('go.xml')
    st=time.time()
    #for i in range(100):
    test(model)
    print(time.time()-st)